import torch

def train_step(model, data_loader, optimizer, loss_fn, device):
    model.train()
    total_loss = 0
    for noisy_images, clean_images in data_loader:
        noisy_images = noisy_images.to(device)
        clean_images = clean_images.to(device)

        optimizer.zero_grad()
        outputs = model(noisy_images)
        loss = loss_fn(outputs, clean_images)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    return total_loss / len(data_loader)

def test_step(model, data_loader, loss_fn, device):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for noisy_images, clean_images in data_loader:
            noisy_images = noisy_images.to(device)
            clean_images = clean_images.to(device)

            outputs = model(noisy_images)
            loss = loss_fn(outputs, clean_images)

            total_loss += loss.item()
    return total_loss / len(data_loader)